"""LightGBM量化模型"""
import pandas as pd
import numpy as np
import lightgbm as lgb
from typing import Dict, Optional
from .base_model import BaseQuantModel


class LightGBMModel(BaseQuantModel):
    """LightGBM模型（支持分类和回归）"""
    
    def __init__(self, task_type: str = "classification", **kwargs):
        """
        初始化LightGBM模型
        
        Args:
            task_type: 任务类型，"classification"或"regression"
            **kwargs: LightGBM参数
                常用参数：
                - n_estimators: 树的数量，默认100
                - max_depth: 树的最大深度，默认-1（无限制）
                - learning_rate: 学习率，默认0.1
                - subsample: 样本采样率，默认1.0
                - colsample_bytree: 特征采样率，默认1.0
        """
        super().__init__(model_name="LightGBM", task_type=task_type)
        
        default_params = {
            'n_estimators': 100,
            'max_depth': -1,
            'learning_rate': 0.1,
            'subsample': 1.0,
            'colsample_bytree': 1.0,
            'random_state': 42,
            'n_jobs': -1,
            'verbosity': -1
        }
        
        if task_type == "classification":
            default_params['objective'] = 'binary'
            default_params['metric'] = 'binary_logloss'
        else:
            default_params['objective'] = 'regression'
            default_params['metric'] = 'rmse'
        
        self.model_params = {**default_params, **kwargs}
    
    def _build_model(self, **kwargs):
        """构建LightGBM模型"""
        params = {**self.model_params, **kwargs}
        
        if self.task_type == "classification":
            self.model = lgb.LGBMClassifier(**params)
        else:
            self.model = lgb.LGBMRegressor(**params)
    
    def _train_model(
        self,
        X_train: pd.DataFrame,
        y_train: pd.Series,
        X_val: Optional[pd.DataFrame] = None,
        y_val: Optional[pd.Series] = None,
        **kwargs
    ) -> Dict:
        """训练LightGBM模型"""
        eval_set = None
        eval_names = None
        
        if X_val is not None and y_val is not None:
            eval_set = [(X_train, y_train), (X_val, y_val)]
            eval_names = ['train', 'val']
        
        self.model.fit(
            X_train,
            y_train,
            eval_set=eval_set,
            eval_names=eval_names,
            verbose=False
        )
        
        # 计算训练集得分
        train_score = self.model.score(X_train, y_train)
        
        result = {
            'train_score': train_score
        }
        
        # 如果有验证集，计算验证集得分
        if X_val is not None and y_val is not None:
            val_score = self.model.score(X_val, y_val)
            result['val_score'] = val_score
        
        # 获取特征重要性
        feature_importance = self.model.feature_importances_
        result['feature_importance'] = dict(zip(X_train.columns, feature_importance))
        
        return result
    
    def _predict_proba(self, X: pd.DataFrame) -> np.ndarray:
        """预测概率或值"""
        if self.task_type == "classification":
            proba = self.model.predict_proba(X)
            # 返回正类概率
            if proba.shape[1] == 2:
                return proba[:, 1]
            return proba
        else:
            return self.model.predict(X)

